#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>

#define NUM_COLORS 3  // 定义颜色数量

int main(int argc, char** argv) {
    int world_rank, world_size;
    int color, key;
    MPI_Comm new_comm;
    
    // 初始化 MPI 环境
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);
    
    printf("进程 %d/%d: 程序启动\n", world_rank, world_size);
    
    // 根据进程排名分配颜色
    printf("演示1: 基本的分裂操作:\n");
    color = world_rank % NUM_COLORS;
    key = world_rank;  // 保持原始排序
    
    printf("进程 %d: 颜色=%d, 键值=%d\n", world_rank, color, key);
    
    // 调用 MPI_Comm_split - 这里会被拦截
    MPI_Comm_split(MPI_COMM_WORLD, color, key, &new_comm);
    
    if (new_comm == MPI_COMM_NULL) {
        printf("进程 %d: 未分配到任何通信器\n", world_rank);
    } else {
        int new_rank, new_size;
        MPI_Comm_rank(new_comm, &new_rank);
        MPI_Comm_size(new_comm, &new_size);
        
        printf("进程 %d: 在新通信器中的排名=%d/%d\n", 
               world_rank, new_rank, new_size);
        
        // 释放通信器
        MPI_Comm_free(&new_comm);
    }

    if (world_rank == 0) {
        printf("\n演示2 MPI_UNDEFINED 的使用:\n");
    }
    MPI_Barrier(MPI_COMM_WORLD);
    
    // 只有排名为偶数的进程参与分裂
    if (world_rank % 2 == 0) {
        color = 1;  // 所有偶数进程使用相同颜色
        key = world_rank;
    } else {
        color = MPI_UNDEFINED;  // 奇数进程不参与
    }
    
    MPI_Comm_split(MPI_COMM_WORLD, color, key, &new_comm);
    
    if (new_comm == MPI_COMM_NULL) {
        printf("进程 %d: 未分配到新通信器 (使用了 MPI_UNDEFINED)\n", world_rank);
    } else {
        int new_rank, new_size;
        MPI_Comm_rank(new_comm, &new_rank);
        MPI_Comm_size(new_comm, &new_size);
        
        printf("进程 %d: 在偶数进程通信器中的排名=%d/%d\n", 
               world_rank, new_rank, new_size);
        
        MPI_Comm_free(&new_comm);
    }
    if (world_rank == 0) {
        printf("\n演示3键值排序:\n");
    }
    MPI_Barrier(MPI_COMM_WORLD);
    
    // 所有进程使用相同颜色，但键值反向排序
    color = 0;
    key = world_size - world_rank - 1;  // 反向键值
    
    MPI_Comm_split(MPI_COMM_WORLD, color, key, &new_comm);
    
    if (new_comm != MPI_COMM_NULL) {
        int new_rank, new_size;
        MPI_Comm_rank(new_comm, &new_rank);
        MPI_Comm_size(new_comm, &new_size);
        
        printf("进程 %d: 原始排名=%d, 键值=%d, 新排名=%d\n", 
               world_rank, world_rank, key, new_rank);
        
        MPI_Comm_free(&new_comm);
    }
    if (world_rank == 0) {
        printf("\n演示4多个通信器创建:\n");
    }
    MPI_Barrier(MPI_COMM_WORLD);
    
    MPI_Comm comm1, comm2;
    
    // 第一次分裂
    color = world_rank / 2;  // 每2个进程一组
    MPI_Comm_split(MPI_COMM_WORLD, color, world_rank, &comm1);
    
    // 第二次分裂
    color = world_rank % 2;  // 奇偶分组
    MPI_Comm_split(MPI_COMM_WORLD, color, world_rank, &comm2);
    
    if (comm1 != MPI_COMM_NULL) {
        int rank1, size1;
        MPI_Comm_rank(comm1, &rank1);
        MPI_Comm_size(comm1, &size1);
        printf("进程 %d: 在通信器1(每2个一组)中的排名=%d/%d\n", world_rank, rank1, size1);
        MPI_Comm_free(&comm1);
    }
    
    if (comm2 != MPI_COMM_NULL) {
        int rank2, size2;
        MPI_Comm_rank(comm2, &rank2);
        MPI_Comm_size(comm2, &size2);
        printf("进程 %d: 在通信器2(奇偶分组)中的排名=%d/%d\n", world_rank, rank2, size2);
        MPI_Comm_free(&comm2);
    }
    printf("进程 %d: 程序结束\n", world_rank);
    MPI_Finalize();
    return 0;
}